import sys
import os
import json
import xml.etree.ElementTree as ET
import shutil

START_BOUNDING_BOX_ID = 0
PRE_DEFINE_CATEGORIES = {"breast": 0}  # 修改的地方，修改为自己的类别
VOC_DIR = "F:\\datasets\\voc\\breast\\VOCdevkit\\VOC2007"
COCO_DIR = "F:\\datasets\\voc\\breast\\COCO"


def get_xml_item_by_name(root, name):
    return root.findall(name)


def get_filename_as_int(filename):
    try:
        filename = os.path.splitext(filename)[0]
        return filename
    except:
        raise NotImplementedError('Filename %s is supposed to be an integer.' % filename)


def convert(type, year):
    itemNameList = open(os.path.join(VOC_DIR, f'ImageSets\\Main\\{type}.txt'), 'r')
    content = {"images": [], "type": "instances", "annotations": [],
               "categories": []}
    categories = PRE_DEFINE_CATEGORIES
    bnd_id = START_BOUNDING_BOX_ID
    for itemName in itemNameList:
        xmlAnnotationFilePath = os.path.join(VOC_DIR, "Annotations", f"{itemName.strip()}.xml")
        tree = ET.parse(xmlAnnotationFilePath)
        root = tree.getroot()
        imageFileName = get_xml_item_by_name(root, 'filename')[0].text
        imageFilePath = get_xml_item_by_name(root, 'path')[0].text
        image_id = get_filename_as_int(imageFileName)
        size = get_xml_item_by_name(root, 'size')
        width = int(get_xml_item_by_name(size[0], "width")[0].text)
        height = int(get_xml_item_by_name(size[0], "height")[0].text)
        image = {'file_name': imageFileName, 'height': height, 'width': width,
                 'id': image_id}
        content['images'].append(image)
        for obj in get_xml_item_by_name(root, 'object'):
            category = get_xml_item_by_name(obj, 'name')[0].text
            if category not in categories:
                new_id = len(categories)
                categories[category] = new_id
            category_id = categories[category]
            bndbox = get_xml_item_by_name(obj, 'bndbox')[0]
            xmin = int(get_xml_item_by_name(bndbox, 'xmin')[0].text) - 1
            ymin = int(get_xml_item_by_name(bndbox, 'ymin')[0].text) - 1
            xmax = int(get_xml_item_by_name(bndbox, 'xmax')[0].text)
            ymax = int(get_xml_item_by_name(bndbox, 'ymax')[0].text)
            assert (xmax > xmin)
            assert (ymax > ymin)
            o_width = abs(xmax - xmin)
            o_height = abs(ymax - ymin)
            ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id':
                image_id, 'bbox': [xmin, ymin, o_width, o_height],
                   'category_id': category_id, 'id': bnd_id, 'ignore': 0,
                   'segmentation': []}
            content['annotations'].append(ann)
            bnd_id = bnd_id + 1
        for cate, cid in categories.items():
            cat = {'supercategory': 'none', 'id': cid, 'name': cate}
            content['categories'].append(cat)
        # 复制相关图片
        shutil.copyfile(imageFilePath, os.path.join(COCO_DIR, f"COCO{year}", f"{type}{year}", imageFileName))
    json_fp = open(os.path.join(COCO_DIR, f"COCO{year}", "annotations", f"instances_{type}{year}.json"), 'w')
    json_str = json.dumps(content)
    json_fp.write(json_str)
    json_fp.close()
    itemNameList.close()


if __name__ == '__main__':
    convert("train", 2017)
    convert("val", 2017)
    convert("test", 2017)
